import socket
import datetime
import sys
import ipaddress
import threading
import os
from scapy.sendrecv import sr, sr1
from scapy.layers.inet import IP, TCP
from utils import createIPList
BLUE, RED, WHITE, YELLOW, MAGENTA, GREEN, END = '', '', '', '', '', '', ''

class ThreadManager(object):
    i = 0

    def __init__(self, ipList):
        self.allIps = ipList
        self.size = len(ipList)

    def getNextIp(self):
        if not (self.i >= self.size - 1):
            ip = self.allIps[self.i]
            self.i += 1
            return ip
        return 0

    def getID(self):
        return self.i + 1

def coreOptions():
    options = [["network", "IP range to scan", ""], ["port-timeout", "Timeout (in sec) for port 80.", "0.3"],
               ["threads", "Number of threads to run.", "50"], ["verbose", "Show verbose output.", "true"]]
    return options


# def createIPList(network):
#     net4 = ipaddress.ip_network(network)
#     ipList = []
#     for x in net4.hosts():
#         ipList.append(x)
#     return ipList

def print1(data):
    global ui
    if verbose:
        print(data)
        ui.logtext.append(data)

# def syn_scan_final(hostname,lport,hport):
#     result_raw = sr(IP(dst = hostname)/TCP(dport = (int(lport),int(hport)),flags=2),timeout = 1,verbose = False)#dport = (int(lport),int(hport))可以扫描指定的端口范围，这也是scapy的一个好处，flags=2表示syn包（三次握手的第一个包）。sr是发送三层数据包,timeout是定义超时时间,verbose如果不将该参数定义为False，那么屏幕将会打印一些不相关的字符串，定义为False之后，就可以打印自己想看到的字符串
#     result_list = result_raw[0].res #表示接收到的数据包的第0层里面的res全部赋值给result_list
#     for i in range(len(result_list)): #主要意思是将收到的每个包的下标赋值给i，如接收到了1024个包，那么i就等于0、1、2、3...（如图一、二）
#         if result_list[i][1].haslayer(TCP): #在i层里的1层中，有没有TCP这一层，如果有执行
#             TCP_Fields = result_list[i][1].getlayer(TCP).fields #在i层里的1层中，将TCP层里的fields里面的值赋值给TCP_Fields
#             if TCP_Fields['flags'] == 18: #这个flags == 18，是指SYN+ACK包（如图三、四）
#                 print('port:' + str(TCP_Fields['sport'])) #打印开放的端口

def checkServer(address, port):
    if method == "connect":
        s = socket.socket()
        s.settimeout(float(portTimeout))
        try:
            s.connect((address, port))
            data = s.recv(4096)
            s.close()
            return ["True", data]
        except socket.error:
            s.close()
            return ["False"]
        except:
            s.close()
            return ["FAIL"]
    elif method == "syn":
        try:
            temp = sr1(IP(dst=address) /
                      TCP(dport=int(port), flags = 'A'), timeout = 0.3)
            # print(type(temp[0][TCP].res))
            print(temp)
            if temp[0].res:
                result = temp[0].res  # temp分回复和无回显
                if (result[0][1].payload.flags) == 'SA':
                    return ["True"]
                else:
                    return ["False"]
            else:
                print('Nothing got')
                return ["False"]
        except Exception as e:
            print(e)
            return ["FAIL"]


def writeToFile(line):
    file = open(fileName, "a")
    file.write(line)
    file.close()


def restart_line():
    sys.stdout.write('\r')
    sys.stdout.flush()


def statusWidget():
    sys.stdout.write(GREEN + "[" + status + "] " + YELLOW + str(threadManager.getID()) + GREEN + " / " + YELLOW + str(
        allIPs) + GREEN + " hosts done." + END)
    restart_line()
    sys.stdout.flush()


def scan(i):
    global status
    global openPorts
    global done

    while True:
        if stop:
            sys.exit()
        ip = threadManager.getNextIp()
        if ip == 0:
            break
        status = (threadManager.getID() / allIPs) * 100
        status = format(round(status, 2))
        status = str(status) + "%"
        stringIP = str(ip)
        if stop:
            sys.exit()

        isUp = checkServer(stringIP, port)
        if isUp[0] != "FAIL":
            if isUp[0] == "True":
                openPorts = openPorts + 1
                if method == 'connect':
                    print1("[+] Port " + str(port) +  " is open on '" + stringIP + "' - Connection response: " + str(isUp[1]))
                    ui.result.append(GREEN + "[+] Port " + str(port) +  " is open on '" + stringIP + "' - Connection response: " + str(isUp[1]))
                else:
                    print1("[+] Port " + str(port) + " is open on '" + stringIP)
                    ui.result.append("[+] Port " + str(port) + " is open on '" + stringIP)
                logLine = stringIP + " - " + str(isUp[1]) + "\n"
                logLines.append(logLine)
            elif not isUp[0] == "True":
                print1(RED + "[-] Port " + str(port) + " is closed on '" + stringIP + "'" + END)
        else:
            print1(RED + "[!] Failed connecting to '" + stringIP + "'" + END)
    done = done + 1


def core(moduleOptions, ori_ui):
    print(
        "\n" + GREEN + "SSH module by @luyi. Scanning subnet '" + YELLOW + moduleOptions[0][2] + GREEN + "'...\n")

    global status
    global fileName
    global allIPs
    global portTimeout
    global ips
    global threadCount
    global done
    global verbose
    global stop
    global port
    global openPorts
    global logLines
    global threadManager
    logLines = []
    stop = False
    done = 0
    global ui
    ui = ori_ui
    global method
    method = 'connect' if ui.connect.isChecked() else 'syn'
    portTimeout = moduleOptions[1][2]
    network = str(moduleOptions[0][2])
    threadCount = int(moduleOptions[2][2])
    verbose = moduleOptions[3][2]

    if verbose == "true":
        verbose = True
    else:
        verbose = False

    try:
        ipList = createIPList(network)
        allIPs = len(ipList)
        if allIPs == 0:
            raise Exception
    except:
        print(RED + "[!] Invalid subnet. Exiting...\n")
        return

    threadManager = ThreadManager(ipList)

    i = datetime.datetime.now()
    i = str(i).replace(" ", "_")
    i = str(i).replace(":", "-")
    script_path = os.path.dirname(os.path.realpath(__file__))
    script_path = script_path.replace("modules", "")
    if not os.path.exists(script_path + "logs"):
        os.makedirs(script_path + "logs")
    fileName = script_path + "logs/log-ssh-portSpider-" + i + ".log"

    file = open(fileName, 'w')
    file.write("subnet: " + network + "\n")
    file.close()

    port = 22

    openPorts = 0
    threads = []
    for i in range(threadCount):
        i -= 1
        t = threading.Thread(target=scan, args=(i,))
        t.daemon = True
        threads.append(t)
        t.start()

    try:
        while True:
            if done == threadCount and threadManager.getID() == allIPs:
                break
            statusWidget()
    except KeyboardInterrupt:
            stop = True
            verbose = False
            print("\n" + RED + "[I] Stopping..." + END)
    stop = True
    verbose = False

    for logLine in logLines:
        try:
            writeToFile(logLine)
        except:
            writeToFile("WRITING-ERROR")

    print("\n\n" + GREEN + "[I] SSH module done. Results saved to '" + YELLOW + fileName + GREEN + "'.\n")
